import openai
from dotenv import load_dotenv
import os

load_dotenv()
openai.api_key = os.getenv('OPENAI_API_KEY')


def generate_story(prompt, length="short"):
    """生成漫画故事情节"""
    try:
        length_map = {
            "short": "约200字的短篇漫画剧情",
            "medium": "约500字的中篇漫画剧情",
            "long": "约1000字的长篇漫画剧情大纲"
        }

        length_desc = length_map.get(length, length_map["short"])

        response = openai.ChatCompletion.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "你是一个专业的漫画编剧。根据用户提示创作有趣的漫画剧情。"},
                {"role": "user", "content": f"创作一个{length_desc}，主题是: {prompt}\n\n"
                                            "请包括以下要素:\n"
                                            "- 主要角色简介\n"
                                            "- 故事开头、发展和高潮\n"
                                            "- 适合漫画化的场景建议"}
            ],
            temperature=0.8,
            max_tokens=2000
        )

        return response.choices[0].message.content.strip()
    except Exception as e:
        print(f"Error generating story: {e}")
        return "无法生成故事，请稍后再试。"